package words;

import java.io.StringReader;
import java.io.BufferedReader;
import java.io.IOException;
import java.util.LinkedList;
import java.util.ListIterator;
import java.util.Iterator;
import java.util.Vector;
import java.util.StringTokenizer;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.PrintWriter;

public class SgmlTokSentDocument {

  TagDefinition m_tagDef;

  // document tokens
  String[] m_vecTokens;
  // POS tags
  String[] m_vecPOSTags;
  // DICT tags
  int[] m_vecDICTags;

  // sentence ending positions
  int[] m_vecESent;
  // annotations list, ordered by position and length
  LinkedList m_lstAnno;

  
  public SgmlTokSentDocument(SgmlDocument doc)
  {
    m_tagDef = doc.m_tagDef;

    // read tokens & set sentence endings
    Vector vecTokens = new Vector();
    Vector vecESent = new Vector();

    try {
      BufferedReader br = new BufferedReader(new StringReader(doc.m_strText));
      String strLine = br.readLine();
      while (strLine != null) {
        StringTokenizer st = new StringTokenizer(strLine);
        if (st.hasMoreTokens()) {
          while (st.hasMoreTokens())
            vecTokens.add(st.nextToken());
          vecESent.add(new Integer(vecTokens.size()));
        }
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.out.println(e);
      System.exit(1);
    }

    m_vecTokens = new String[vecTokens.size()];
    vecTokens.toArray(m_vecTokens);
    
    m_vecESent = new int[vecESent.size()];
    for (int i = 0; i < vecESent.size(); i++)
      m_vecESent[i] = ((Integer) vecESent.get(i)).intValue();

    // synchronize annotations with tokens
    m_lstAnno = new LinkedList();
    int nTight = 0;
    int nTok = 0;
    ListIterator it = doc.m_lstAnno.listIterator();
    while (it.hasNext()) {
      Annotation anno = (Annotation) it.next();
      Annotation newAnno = new Annotation();
      newAnno.m_nId = anno.m_nId;

      // goes forth to find annotation beginning
      while (nTight < anno.m_nBegin)
        nTight += m_vecTokens[nTok++].length();
      newAnno.m_nBegin = nTok;

      int nLen = 0;
      int nT = nTok;
      while (nLen < anno.m_nLen)
        nLen += m_vecTokens[nT++].length();
      newAnno.m_nLen = nT - nTok;

      newAnno.m_mapAtt.putAll(anno.m_mapAtt);

      m_lstAnno.add(newAnno);
    }

    m_vecPOSTags = null;
    m_vecDICTags = null;
  }


  public boolean loadPOSTags(String strFile)
  {
    Vector vecTags = new Vector();
    try {
      BufferedReader br = new BufferedReader(new FileReader(strFile));
      String strLine = br.readLine();
      while (strLine != null) {
        StringTokenizer st = new StringTokenizer(strLine);
        while (st.hasMoreTokens()) {
          String strToken = st.nextToken();
          int nPos = strToken.lastIndexOf('/');
          vecTags.add(strToken.substring(nPos + 1));
        }
       
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.err.println(e);
      return false;
    }

    assert vecTags.size() == m_vecTokens.length;

    m_vecPOSTags = new String[vecTags.size()];
    vecTags.toArray(m_vecPOSTags);

    return true;
  }


  public boolean loadDICTags(String strFile, String strTag)
  {
    m_vecDICTags = new int[m_vecTokens.length];

    SgmlDocument sgmlDoc = new SgmlDocument(m_tagDef);
    sgmlDoc.load(strFile);
    SgmlTokSentDocument sgmlTSDoc = new SgmlTokSentDocument(sgmlDoc);

    int nLast = 0;
    ListIterator it = sgmlTSDoc.m_lstAnno.listIterator();
    while (it.hasNext()) {
      Annotation anno = (Annotation) it.next();
      if (strTag.compareTo(m_tagDef.getName(anno.m_nId)) == 0) {
        double conf = Double.parseDouble((String) anno.m_mapAtt.get("conf"));
        if (conf >= 0.5) {
          for (int i = nLast; i < anno.m_nBegin; i++)
            m_vecDICTags[i] = 0; // outside NE
	  if (m_vecTokens.length <= anno.m_nBegin) {
	      System.out.println("# of tokens: " + m_vecTokens.length + " in " + strFile);
	      System.out.println(anno.m_nBegin);
	  }
          if (anno.m_nLen == 1)
            m_vecDICTags[anno.m_nBegin] = 4; // unique NE
          else {
            m_vecDICTags[anno.m_nBegin] = 1; // begin NE
            for (int i = anno.m_nBegin + 1; i < anno.m_nBegin + anno.m_nLen - 1; i++)
              m_vecDICTags[i] = 2; // inside NE
            m_vecDICTags[anno.m_nBegin + anno.m_nLen - 1] = 3; // end NE
          }
          nLast = anno.m_nBegin + anno.m_nLen;
        }
      }
    }
    // tag the remaining tokens
    for (int i = nLast; i < m_vecDICTags.length; i++)
      m_vecDICTags[i] = 0; // outside NE

    return true;
  }


  public String toString()
  {
    // initialize out vector
    LinkedList[] vecOut = new LinkedList[m_vecTokens.length];
    for (int i = 0; i < vecOut.length; i++) {
      vecOut[i] = new LinkedList();
      vecOut[i].add(m_vecTokens[i]);
    }

    // augment out vector with annotations
    ListIterator it = m_lstAnno.listIterator();
    while (it.hasNext()) {
      Annotation anno = (Annotation) it.next();
      vecOut[anno.m_nBegin].addFirst(anno.getBegin(m_tagDef));
      vecOut[anno.m_nBegin + anno.m_nLen - 1].addLast(anno.getEnd(m_tagDef));
    }

    // create output string
    StringBuffer sbOut = new StringBuffer();
    int j = 0;
    for (int i = 0; i < m_vecESent.length; i++) {
      for (; j < m_vecESent[i]; j++) {
        ListIterator lit = vecOut[j].listIterator();
        while (lit.hasNext()) {
          sbOut.append(lit.next());
          sbOut.append(" ");
        }
      }
      sbOut.append("\n");
    }

    return sbOut.toString();
  }


  public String toStringFlat(String strTag)
  {
    // initialize out vector
    LinkedList[] vecOut = new LinkedList[m_vecTokens.length];
    for (int i = 0; i < vecOut.length; i++) {
      vecOut[i] = new LinkedList();
      vecOut[i].add(m_vecTokens[i]);
    }

    // augment out vector with annotations, flat structure for 'strTag'
    ListIterator it = m_lstAnno.listIterator();
    int nLast = 0;
    while (it.hasNext()) {
      Annotation anno = (Annotation) it.next();
      String tag = m_tagDef.getName(anno.m_nId);
      if (tag.compareTo(strTag) == 0) {
        if (anno.m_nBegin >= nLast) {
          vecOut[anno.m_nBegin].addFirst(anno.getBegin(m_tagDef));
          vecOut[anno.m_nBegin + anno.m_nLen - 1].addLast(anno.getEnd(m_tagDef));
          nLast = anno.m_nBegin + anno.m_nLen;
        }
      }
      else {
        vecOut[anno.m_nBegin].addFirst(anno.getBegin(m_tagDef));
        vecOut[anno.m_nBegin + anno.m_nLen - 1].addLast(anno.getEnd(m_tagDef));
      }
    }

    // create output string
    StringBuffer sbOut = new StringBuffer();
    int j = 0;
    for (int i = 0; i < m_vecESent.length; i++) {
      for (; j < m_vecESent[i]; j++) {
        ListIterator lit = vecOut[j].listIterator();
        while (lit.hasNext()) {
          sbOut.append(lit.next());
          sbOut.append(" ");
        }
      }
      sbOut.append("\n");
    }

    return sbOut.toString();
  }


  public SgmlTokSentence[] getSentences()
  {
    SgmlTokSentence[] vecSent = new SgmlTokSentence[m_vecESent.length];

    int j = 0;
    ListIterator it = m_lstAnno.listIterator();
    for (int i = 0; i < m_vecESent.length; i++) {
      String[] vecTokens = new String[m_vecESent[i] - j];
      int first = j;
      while (j < m_vecESent[i]) {
        vecTokens[j - first] = m_vecTokens[j];
        j++;
      }

      LinkedList lstAnno = new LinkedList();
      while (it.hasNext()) {
        Annotation anno = (Annotation) it.next();
        if (anno.m_nBegin >= m_vecESent[i]) {
          if (it.hasPrevious())
            it.previous();
          break;
        }

        if (anno.m_nBegin + anno.m_nLen < m_vecESent[i]) {
          Annotation newAnno = new Annotation(anno.m_nBegin - first, anno.m_nLen);
          newAnno.m_nId = anno.m_nId;
          newAnno.m_mapAtt.putAll(anno.m_mapAtt);
          lstAnno.add(newAnno);
        }
      }
      vecSent[i] = new SgmlTokSentence(i, m_tagDef, vecTokens, lstAnno);

      // set sentence POS tags
      if (m_vecPOSTags != null) {
        String[] vecPOSTags = new String[vecTokens.length];
        int k = first;
        while (k < m_vecESent[i]) {
          vecPOSTags[k - first] = new String(m_vecPOSTags[k]);
          k++;
        }
        vecSent[i].setPOSTags(vecPOSTags);
      }

      // set sentence DICT tags
      if (m_vecDICTags != null) {
        int[] vecDICTags = new int[vecTokens.length];
        int k = first;
        while (k < m_vecESent[i]) {
          vecDICTags[k - first] = m_vecDICTags[k];
          k++;
        }
        vecSent[i].setDICTags(vecDICTags);
      }
    }

    return vecSent;
  }


  public int size()
  {
    return m_vecTokens.length;
  }


  static public void main(String[] args) 
  {
    if (args.length != 4) {
      System.out.println("Usage: java SgmlTokSentDocument <corpus-list> <tag-list> " +
                         "<input-dir> <output-dir>!");
      System.exit(1);
    }

    TagDefinition tagDef = new TagDefinition();
    tagDef.load(args[1]);

    try {
      BufferedReader br = new BufferedReader(new FileReader(args[0]));
      String strLine = br.readLine();
      while (strLine != null) {
        strLine = strLine.trim();
        if (strLine.length() > 0) {
          // load input document
          String strInFile = args[2] + "/" + strLine;

          System.out.print("processing [" + strInFile + "] ... ");

          SgmlDocument sgmlDoc = new SgmlDocument(tagDef);
          sgmlDoc.load(strInFile);
          SgmlTokSentDocument sgmlTSDoc = new SgmlTokSentDocument(sgmlDoc);
	  
          // open output document
          String strOutFile = args[3] + "/" + strLine;
          PrintWriter pw = new PrintWriter(new FileWriter(strOutFile));

          pw.println(sgmlTSDoc.toStringFlat("prot"));
          pw.close();

          System.out.println("done!");
        }
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.err.println(e);
      System.exit(1);
    }
  }
  
}
